package com.treasure.infrastructure.repository.extraction.elastic;

import com.treasure.infrastructure.model.doc.SegmentESDocBean;
import com.treasure.infrastructure.repository.extraction.elastic.SegmentESDocRepository;
import lombok.extern.slf4j.Slf4j;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.MatchQueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.search.sort.SortOrder;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.domain.PageRequest;
import org.springframework.data.elasticsearch.core.ElasticsearchRestTemplate;
import org.springframework.data.elasticsearch.core.SearchHit;
import org.springframework.data.elasticsearch.core.SearchHits;
import org.springframework.data.elasticsearch.core.mapping.IndexCoordinates;
import org.springframework.data.elasticsearch.core.query.*;
import org.springframework.stereotype.Service;
import org.springframework.util.CollectionUtils;

import java.util.ArrayList;
import java.util.List;

import static org.elasticsearch.search.sort.SortBuilders.fieldSort;

@Service
@Slf4j
public class SegmentESDocRepositoryImpl implements SegmentESDocRepository {

    @Autowired
    private ElasticsearchRestTemplate elasticsearchRestTemplate;

    @Override
    public void indexSegmentDoc(List<SegmentESDocBean> segmentDocBeanList) {
        if (CollectionUtils.isEmpty(segmentDocBeanList)) {
            return;
        }

        IndexCoordinates indexCoordinates = IndexCoordinates.of(SegmentESDocBean.INDEX_NAME);
        for (SegmentESDocBean segmentDocBean : segmentDocBeanList) {
            IndexQuery indexQueryBuilder = new IndexQueryBuilder()
                    .withId(String.valueOf(System.nanoTime()))
                    .withObject(segmentDocBean)
                    .build();

            String index = elasticsearchRestTemplate.index(indexQueryBuilder, indexCoordinates);
            log.info("indexDoc: " + index);
        }
    }

    @Override
    public List<SegmentESDocBean> searchSegmentDoc(String queryText) {
        IndexCoordinates indexCoordinates = IndexCoordinates.of(SegmentESDocBean.INDEX_NAME);
        BoolQueryBuilder boolQueryBuilder = QueryBuilders.boolQuery();

        MatchQueryBuilder passages = QueryBuilders.matchQuery("passages", queryText);
        boolQueryBuilder.must(passages);

        MatchQueryBuilder docName = QueryBuilders.matchQuery("docName", queryText);
        MatchQueryBuilder summaries = QueryBuilders.matchQuery("summaries", queryText);
        boolQueryBuilder.should(docName).should(summaries);

        NativeSearchQueryBuilder searchQueryBuilder = new NativeSearchQueryBuilder()
                .withQuery(boolQueryBuilder)
                .withSort(fieldSort("createTime").order(SortOrder.DESC));
        NativeSearchQuery query = searchQueryBuilder.build();
        query.setPageable(PageRequest.of(0,1));

        SearchHits<SegmentESDocBean> search = elasticsearchRestTemplate.search(query, SegmentESDocBean.class, indexCoordinates);
        List<SearchHit<SegmentESDocBean>> searchHits = search.getSearchHits();
        List<SegmentESDocBean> result = new ArrayList<>(searchHits.size());
        for (SearchHit<SegmentESDocBean> searchHit : searchHits) {
            SegmentESDocBean content = searchHit.getContent();
            result.add(content);
        }

        return result;
    }


}
